from spm_ADEM import spm_vec
import numpy as np


def isvec(x):
    if isinstance(x,np.ndarray):
        if len(x.shape) == 1:
            return True
        elif len(x.shape) == 2 and min(x.shape) == 1:
            return True
    return False

def spm_diff(f=None,n=[],dx=np.exp(-8),x=None,V=None,q=True,reference=None):
    ''' returns df/dx[0], d^2f/dx[n[0]dx[x[1]],d^3f/dx[n[0]]dx[n[1]]dx[n[2]],etc.'''
    if f is None:
        raise TypeError('Please specify function f')
    if x is None:
        x = []
    if dx < 0:
        raise TypeError('dx should be positive')
    if V is None:
        V = [[]] * len(n)
    
    # The underlying reason is that reasond is that the gradient taken is
    # f(x+d)-f(x) so that dt should be negative to offset the wrong sign
    dx = -dx 
    xv = [spm_vec(xi) for xi in x]
    
    # Prepare matrixes    
    '''
    if len(V) < len(n):
        V += [[]]*(len(n)-len(V))
    
    for i in xrange(len(V)):
        if V[i] == []:
            pass
        elif isinstance(V[i],np.ndarray):
            if max(V[i].shape) < 1:
                V[i] = []
            elif max(V[i].shape) == 1:
                V[i] = np.ones(max(x[i].shape))
            elif max(V[i].shape) != max(x[i].shape):
                raise TypeError
            elif np.prod(V[i]) == x[i].shape:
                V[i] = V[i].flat
            elif V[i].shape[0] == V[i].shape[1]:
                pass
    '''
    y1,y0 = _num_diff(f,n=n,dx=dx,x=xv,V=V,q=q)
    return y1,y0
       
def _num_diff(f,n,dx,x,V,q):
    if len(n) == 1:    
        # Unpack arguments and evaluate function
        # ---------------------------------------------------------------------
        xi   = [i.data for i in x]
        dfdx = [[]]*len(x[n[0]])        
        for i in xrange(len(x[n[0]])):
            y0 = f(*xi)
            try:
                y0.vec
            except AttributeError:
                y0 = spm_vec(y0)
            x[n[0]].vec[i] += dx
            y1             = f(*xi)
            try:
                y1.vec
            except AttributeError:
                y1 = spm_vec(y1)
            x[n[0]].vec[i]    -= dx
            # Finite difference
            # ---------------------------------------------------------------------
            y1 -= y0
            y1 /= dx
            dfdx[i] = y1
            
        if q and isvec(x[n[0]].data):
            # Write in c ?
            J = np.empty((len(x[n[0]]),len(y1)),dtype='d')
            for i in xrange(len(dfdx)):
                J[i,:] = dfdx[i].vec
            return [spm_vec(J)],y0
        return [spm_vec(dfdx).vec],y0
        
        
    elif len(n) > 1:
        def __num_diff(xt):
            y1,y0 = _num_diff(f,n=n[0:-1],dx=dx,x=xt,V=V,q=q)
            return y1,y0
        # Unpack arguments and evaluate function
        # ---------------------------------------------------------------------
        m = len(x[n[-1]])
        dy0,y    = __num_diff(x)
        dfdx     = [[]] * m
        for i in xrange(m):
            x[n[-1]].vec[i] += dx
            dy1      = __num_diff(x)[0]
            x[n[-1]].vec[i] -= dx
            dy1[0]   -= dy0[0]
            dy1[0]   /= dx
            dfdx[i]  = dy1[0]
        return [spm_vec(dfdx)]+dy0,y
    else:
        return [],f(*[i.data for i in x])  

if __name__ == '__main__':
    def f(x,y,v):
        return np.array([(3.*x[0]-x[1]**3.*2.+2.*y+4.*v*x[2]),(3.*x[0]*x[2]-x[1]**2.*2.+2.*y+4.*v)])
    r = spm_num_diff(f=f,n=[0,0],dx=np.exp(-8,dtype='d'),x=[np.array([1.,3.,4.]),np.array(2.),np.array(3.)])
    import time
    t = time.time()
    for i in xrange(40000): spm_num_diff(f=f,n=[0,0],dx=np.exp(-8,dtype='d'),x=[np.array([1.,3.,4.]),np.array(2.),np.array(3.)])
    print time.time() - t
